import pandas as pd
from sklearn.cluster import KMeans
from sklearn.preprocessing import StandardScaler
import matplotlib.pyplot as plt
import seaborn as sns
import warnings
warnings.filterwarnings('ignore')

class ClusterUtils(object):
    def __init__(self):
        self.df = pd.read_csv('./scenic_data.csv')
    
    def get_cluster(self):
        # 提取特征并标准化
        features = self.df[['non_weekend', 'out_province_ratio', 'elderly_ratio']]
        scaler = StandardScaler()
        scaler_features = scaler.fit_transform(features)

        # 使用K-means聚类
        kmeans = KMeans(n_clusters=4, random_state=42)
        self.df['cluster'] = kmeans.fit_predict(scaler_features)
        
        return self.df
    
    def visualize_clusters(self):
        # 可视化聚类结果
        plt.figure(figsize=(10, 6))
        
        # 使用散点图展示聚类结果，以两个主要特征为例
        sns.scatterplot(
            x='out_province_ratio', 
            y='elderly_ratio',
            hue='cluster',
            palette=['red', 'green', 'blue', 'purple'],
            data=self.df,
            s=100,
            alpha=0.7
        )
        
        plt.title('K-means Clustering Results (k=4)')
        plt.xlabel('Out of Province Visitor Ratio')
        plt.ylabel('Elderly Visitor Ratio')
        plt.grid(True)
        plt.legend(title='Cluster')
        plt.tight_layout()
        
        # 保存图像
        plt.savefig('cluster_visualization.png', dpi=300)
        plt.show()
        
        # 可选：输出聚类中心
        print("Cluster Centers:")
        print(self.df.groupby('cluster').mean())
        
        # 保存聚类结果到新文件
        self.df.to_csv('clustered_scenic_data.csv', index=False)

if __name__ == '__main__':
    cluster_util = ClusterUtils()
    clustered_data = cluster_util.get_cluster()
    cluster_util.visualize_clusters()